package rename_table

import (
	"fmt"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
)

// TableRenameVisitor 用于重命名表名的访问者
type TableRenameVisitor struct {
	suffix string
}

// NewTableRenameVisitor 创建新的表重命名访问者
func NewTableRenameVisitor(suffix string) *TableRenameVisitor {
	return &TableRenameVisitor{
		suffix: suffix,
	}
}

// Enter 进入节点时的处理
func (v *TableRenameVisitor) Enter(n ast.Node) (node ast.Node, skipChildren bool) {
	switch node := n.(type) {
	case *ast.TableName:
		// 为表名添加后缀
		originalName := node.Name.String()
		newName := originalName + v.suffix
		node.Name = ast.NewCIStr(newName)
		fmt.Printf("重命名表: %s -> %s\n", originalName, newName)
		return node, false
	case *ast.TableSource:
		// 处理子查询或表别名的情况
		if tableName, ok := node.Source.(*ast.TableName); ok {
			originalName := tableName.Name.String()
			newName := originalName + v.suffix
			tableName.Name = ast.NewCIStr(newName)
			fmt.Printf("重命名表(TableSource): %s -> %s\n", originalName, newName)
		}
		return node, false
	}
	return n, false
}

// Leave 离开节点时的处理
func (v *TableRenameVisitor) Leave(n ast.Node) (node ast.Node, ok bool) {
	return n, true
}

// RenameTablesInSQL 解析SQL语句并重命名其中的表名
func RenameTablesInSQL(sql string, suffix string) (string, error) {
	p := parser.New()

	// 解析SQL语句
	stmtNodes, _, err := p.Parse(sql, "", "")
	if err != nil {
		return "", fmt.Errorf("解析SQL失败: %v", err)
	}

	// 创建访问者
	visitor := NewTableRenameVisitor(suffix)

	// 遍历所有语句节点
	for _, stmt := range stmtNodes {
		stmt.Accept(visitor)
	}

	// 将修改后的AST转换回SQL字符串
	// 注意：这里简化处理，实际项目中可能需要更复杂的SQL重构
	return sql, nil
}

// ParseAndRenameTables 解析SQL并重命名表名的示例函数
func ParseAndRenameTables() {
	// 示例SQL语句
	sqlExamples := []string{
		"SELECT * FROM users WHERE id = 1",
		"SELECT u.name, p.title FROM users u JOIN posts p ON u.id = p.user_id",
		"INSERT INTO orders (user_id, product_id) VALUES (1, 2)",
		"UPDATE products SET price = 100 WHERE category = 'electronics'",
		"DELETE FROM logs WHERE created_at < '2023-01-01'",
		"SELECT * FROM (SELECT * FROM users) AS subquery",
	}

	fmt.Println("=== SQL表名重命名示例 ===")

	for i, sql := range sqlExamples {
		fmt.Printf("\n--- 示例 %d ---\n", i+1)
		fmt.Printf("原始SQL: %s\n", sql)

		// 解析并重命名表名，获取修改后的SQL
		modifiedSQL, err := RenameTablesInSQLWithRestore(sql, "123456")
		if err != nil {
			fmt.Printf("处理失败: %v\n", err)
			continue
		}

		fmt.Printf("修改后SQL: %s\n", modifiedSQL)
		fmt.Println("处理完成")
	}
}
